package com.taint.agent;

import com.taint.callgraph.TraceItem;
import com.taint.callgraph.CallRecorder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.loading.ClassInjector;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class FlowTaintClassLoader {


    /*
    https://github.com/signalfx/signalfx-java-tracing/blob/113c952dd180413d2d51ba97ff97cf13a61115ed/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/SpockRunner.java#L156
     */
    private static JarFile agentFile ;

    public static final String[] BOOTSTRAP_PACKAGE_PREFIXES = {

            "org.apache.log4j",
            //"org.yaml.snakeyaml",
            "com.taint.logger",
            "com.taint.agent",
            "com.taint.configuration",
            "com.taint.sourcesink",
            "com.taint.callgraph",
    };

    private static boolean isNeedBootstrap(String className) {
        for (int i = 0; i < BOOTSTRAP_PACKAGE_PREFIXES.length; ++i) {
            if (className.startsWith(BOOTSTRAP_PACKAGE_PREFIXES[i])) {
                return true;
            }
        }
        return false;
    }

    private static String findAgentPath() {
        String jarPath;
        String classResourcePath = FlowTaintAgent.class.getName().replaceAll("\\.", "/") + ".class";
        URL resource = ClassLoader.getSystemClassLoader().getResource(classResourcePath);
        if (resource != null) {
            String urlString = resource.toString();
            if (!urlString.contains("!")) {
                jarPath = urlString.substring(urlString.indexOf("file:"), urlString.indexOf("com"));
            } else {
                jarPath = urlString.substring(urlString.indexOf("file:"), urlString.indexOf("!"));

            }
            File agentJarFile;
            try {
                agentJarFile = new File(new URL(jarPath).toURI());


            } catch (MalformedURLException | URISyntaxException e) {
                throw new RuntimeException("agent jar path can`t find", e);
            }

            return agentJarFile.getParent();
        }
        throw new RuntimeException("agent jar path can`t find");
    }

    public static ArrayList LoadBootstapClass() {
        ArrayList needLoadClassed = new ArrayList<>();

        try {
            agentFile = new JarFile(findAgentPath() + "/Agent.jar");
            //agentFilePath=agentFile.getName();
            Enumeration<JarEntry> entries = agentFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                if (!entry.getName().endsWith(".class")) continue;
                String replace = entry.getName().replace("/", ".");
                if (isNeedBootstrap(replace)) {
                    needLoadClassed.add(replace);
                }
            }


        } catch (IOException e) {
            e.printStackTrace();
        }

        return needLoadClassed;

    }

    public static void injectBootstrapClasses(Instrumentation instrumentation) throws
            IOException, ClassNotFoundException {


        File temp = Files.createTempDirectory("tmp").toFile();


        temp.deleteOnExit();


        ArrayList needLoadClazz = LoadBootstapClass();


        Map<TypeDescription, byte[]> map = new HashMap<>();

        for (int j = 0; j < needLoadClazz.size(); j++) {
            String clazzName = needLoadClazz.get(j).toString().replaceAll("/", ".").replace(".class", "");
            boolean isLoaded = isClassLoaded(clazzName, FlowTaintAgent.class.getClassLoader());
            if (isLoaded) {
                Class<?> clazz = FlowTaintAgent.class.getClassLoader().loadClass(clazzName);
                map.put(new TypeDescription.ForLoadedType(clazz), ClassFileLocator.ForClassLoader.read(clazz));
            }else {
                Class<?>[] loadedClazz = new Class<?>[]{
                        CallRecorder.class,
                        TraceItem.class,
                        com.taint.callgraph.CallGraph.class,
                        com.taint.callgraph.CallStack.class,
                        com.taint.logger.TraceEndpoint.class,
                        com.taint.logger.TraceLog.class,
                        com.taint.logger.TraceEntry.class,
                        com.taint.utils.Utils.class,
                        com.taint.logger.TraceLogShutdown.class,
                        com.taint.logger.TraceLogExceptionHandler.class,
                };
                for (int m = 0; m < loadedClazz.length; m++) {
                    map.put(new TypeDescription.ForLoadedType(loadedClazz[m]), ClassFileLocator.ForClassLoader.read(loadedClazz[m]));
                }
                //map.put(new TypeDescription.ForLoadedType(), ClassFileLocator.ForClassLoader.read(CallItem.class));
            }

        }
        // map.put(new TypeDescription.ForLoadedType(com.taint.agent.TaintFlowAgent.class), ClassFileLocator.ForClassLoader.read(com.taint.agent.TaintFlowAgent.class));

        ClassInjector.UsingInstrumentation.of(temp, ClassInjector.UsingInstrumentation.Target.BOOTSTRAP, instrumentation).inject(map);

        //instrumentation.appendToBootstrapClassLoaderSearch(agentFile);

    }


    public static boolean isClassLoaded(final String className, final ClassLoader classLoader) {
        try {
            final Method findLoadedClassMethod =
                    ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
            try {
                findLoadedClassMethod.setAccessible(true);
                final Class<?> loadedClass =
                        (Class<?>) findLoadedClassMethod.invoke(classLoader, className);
                return null != loadedClass && loadedClass.getClassLoader() == classLoader;
            } catch (final Exception e) {
                throw new IllegalStateException(e);
            } finally {
                findLoadedClassMethod.setAccessible(false);
            }
        } catch (final NoSuchMethodException e) {
            throw new IllegalStateException(e);
        }
    }


}